//
// Created by tang on 22-9-12.
//

#include "vulkan_triangle.hpp"
#include "shaders.hpp"
#include <chrono>
#include <glm/gtc/matrix_transform.hpp>

struct VertexDataStruct {
    glm::vec2 pos;
    glm::vec3 color;

    static vk::VertexInputBindingDescription getBindingDescription()
    {
        vk::VertexInputBindingDescription bindingDescription {};
        bindingDescription.binding   = 0;
        bindingDescription.stride    = sizeof(VertexDataStruct);
        bindingDescription.inputRate = vk::VertexInputRate::eVertex;

        return bindingDescription;
    }

    static std::array<vk::VertexInputAttributeDescription, 2>
    getAttributeDescriptions()
    {
        std::array<vk::VertexInputAttributeDescription, 2>
            attributeDescriptions {};

        attributeDescriptions[0].binding  = 0;
        attributeDescriptions[0].location = 0;
        attributeDescriptions[0].format   = vk::Format::eR32G32Sfloat;
        attributeDescriptions[0].offset   = offsetof(VertexDataStruct, pos);

        attributeDescriptions[1].binding  = 0;
        attributeDescriptions[1].location = 1;
        attributeDescriptions[1].format   = vk::Format::eR32G32B32Sfloat;
        attributeDescriptions[1].offset   = offsetof(VertexDataStruct, color);

        return attributeDescriptions;
    }
};

static const VertexDataStruct vertices[] {
    {{
         -0.5f,
         -0.5f,
     },
     {1.f, 0.f, 0.f}},  // #0 top-left
    {{
         0.5f,
         -0.5f,
     },
     {0.f, 1.f, 0.f}},  // #1 top-right
    {{
         -0.5f,
         0.5f,
     },
     {0.f, 0.f, 1.f}},  // #2 bottom-left
    {{
         0.5f,
         0.5f,
     },
     {1.f, 0.f, 1.f}},  // #3 bottom-right
    {{
         0.0f,
         0.0f,
     },
     {1.f, 1.f, 1.f}},  // #4 center
};

static const u32 vertexIndices[] {4,2,0,4,0,1,4,3,2,4,1,3};

struct UniformBufferObject {
    glm::mat4 model;
    glm::mat4 view;
    glm::mat4 proj;
};

namespace core::app {
    VulkanTriangle::VulkanTriangle() {}

    VulkanTriangle::~VulkanTriangle() {}

    void VulkanTriangle::OnInitialize()
    {
        using namespace vk;
        // create shader modules
        {
            ShaderModuleCreateInfo vertInfo {
                .codeSize = shaders::triangle_vert_length,
                .pCode    = (u32 *)shaders::triangle_vert_data,
            };
            m_vert_shader = m_device->createShaderModuleUnique(vertInfo);
            ShaderModuleCreateInfo fragInfo {
                .codeSize = shaders::triangle_frag_length,
                .pCode    = (u32 *)shaders::triangle_frag_data,
            };
            m_frag_shader = m_device->createShaderModuleUnique(fragInfo);
        }

        // vertex buffer staging
        {
            auto [buffer, mem] =
                CreateBuffer(sizeof(vertices),
                             BufferUsageFlagBits::eTransferSrc,
                             MemoryPropertyFlagBits::eHostVisible);
            UploadBufferData(mem.get(), vertices, sizeof(vertices));
            std::tie(m_vertex_buffer, m_vertex_memory) = CreateBuffer(
                sizeof(vertices), BufferUsageFlagBits::eVertexBuffer | BufferUsageFlagBits::eTransferDst,
                MemoryPropertyFlagBits::eHostVisible);
            CopyBuffer(buffer.get(), m_vertex_buffer.get(), sizeof(vertices));
        }

        // index buffer staging
        {
            auto [buffer, mem] = CreateBuffer(
                sizeof(vertexIndices), BufferUsageFlagBits::eTransferSrc,
                MemoryPropertyFlagBits::eHostVisible);
            UploadBufferData(mem.get(), vertexIndices, sizeof(vertexIndices));
            std::tie(m_index_buffer, m_index_memory) = CreateBuffer(
                sizeof(vertexIndices), BufferUsageFlagBits::eIndexBuffer | BufferUsageFlagBits::eTransferDst,
                MemoryPropertyFlagBits::eHostVisible);
            CopyBuffer(buffer.get(), m_index_buffer.get(), sizeof(vertexIndices));
        }

        // uniform buffer
        {
            m_uniform_buffers.reserve(MaxInFlightImageCount());
            for (u32 i=0; i<MaxInFlightImageCount(); ++i) {
                m_uniform_buffers.emplace_back(BufferMemory {
                    CreateBuffer(sizeof(UniformBufferObject),
                                 BufferUsageFlagBits::eUniformBuffer,
                                 MemoryPropertyFlagBits::eHostVisible)}
                    );
            }
        }

        // descriptor set
        {
            std::vector<DescriptorSetLayout> layouts(MaxInFlightImageCount(), m_descriptor_set_layout.get());
            DescriptorSetAllocateInfo dsai{
                .descriptorPool = m_descriptor_pool.get(),
                .descriptorSetCount = MaxInFlightImageCount(),
                .pSetLayouts = layouts.data(),
            };
            m_uniform_descriptor_sets = m_device->allocateDescriptorSets(dsai);

            for (u32 i = 0; i<MaxInFlightImageCount(); ++i) {
                DescriptorBufferInfo dbi{
                    .buffer = std::get<0>(m_uniform_buffers[i]).get(),
                    .offset = 0,
                    .range = sizeof(UniformBufferObject),
                };

                WriteDescriptorSet wds{
                    .dstSet = m_uniform_descriptor_sets[i],
                    .dstBinding = 0,
                    .dstArrayElement = 0,
                    .descriptorCount = 1,
                    .descriptorType = DescriptorType::eUniformBuffer,
                    .pBufferInfo = &dbi,
                };

                m_device->updateDescriptorSets(wds, {});
            }
        }
    }

    void VulkanTriangle::OnCleanup() {}

    void VulkanTriangle::OnWindowResized() { ReCreateSwapchianAndStuff(); }
    vk::Pipeline
    VulkanTriangle::OnCreateGraphicsPipeline(vk::Extent2D screen_size)
    {
        using namespace vk;

        auto bindingDescription = VertexDataStruct::getBindingDescription();
        auto attributeDescription =
            VertexDataStruct::getAttributeDescriptions();

        PipelineVertexInputStateCreateInfo pcisci {
            .vertexBindingDescriptionCount   = 1,
            .pVertexBindingDescriptions      = &bindingDescription,
            .vertexAttributeDescriptionCount = (u32)attributeDescription.size(),
            .pVertexAttributeDescriptions    = attributeDescription.data()};

        PipelineInputAssemblyStateCreateInfo piasci {
            .topology               = PrimitiveTopology::eTriangleList,
            .primitiveRestartEnable = false};

        Viewport vp {.x        = 0,
                     .y        = 0,
                     .width    = (f32)screen_size.width,
                     .height   = (f32)screen_size.height,
                     .minDepth = 0,
                     .maxDepth = 1};

        Rect2D scissor {.offset = {0, 0},
                        .extent = screen_size};

        PipelineViewportStateCreateInfo pvsci {
            .viewportCount = 1,
            .pViewports    = &vp,
            .scissorCount  = 1,
            .pScissors     = &scissor,
        };

        PipelineRasterizationStateCreateInfo prsci {
            .depthClampEnable        = false,
            .rasterizerDiscardEnable = false,
            .cullMode                = CullModeFlagBits::eBack,
            .frontFace               = FrontFace::eCounterClockwise,
            .depthBiasEnable         = false,
            .depthBiasConstantFactor = 0.f,
            .depthBiasClamp          = 0.f,
            .depthBiasSlopeFactor    = 0.f,
            .lineWidth               = 1.f,
        };

        PipelineMultisampleStateCreateInfo pmsci {
            .rasterizationSamples  = SampleCountFlagBits::e1,
            .sampleShadingEnable   = false,
            .minSampleShading      = 1.f,
            .pSampleMask           = nullptr,
            .alphaToCoverageEnable = false,
            .alphaToOneEnable      = false,
        };

        PipelineColorBlendAttachmentState pcbas {
            .blendEnable         = false,
            .srcColorBlendFactor = BlendFactor::eOne,
            .dstColorBlendFactor = BlendFactor::eZero,
            .colorBlendOp        = BlendOp::eAdd,
            .srcAlphaBlendFactor = BlendFactor::eOne,
            .dstAlphaBlendFactor = BlendFactor::eZero,
            .alphaBlendOp        = BlendOp::eAdd,
            .colorWriteMask =
                ColorComponentFlagBits::eR | ColorComponentFlagBits::eG |
                ColorComponentFlagBits::eB | ColorComponentFlagBits::eA,
        };

        PipelineColorBlendStateCreateInfo pcbsci {
            .logicOpEnable   = false,
            .logicOp         = LogicOp::eCopy,
            .attachmentCount = 1,
            .pAttachments    = &pcbas,
            .blendConstants  = std::array<float, 4> {0, 0, 0, 0}};

        DynamicState states[] {DynamicState::eViewport,
                               DynamicState::eLineWidth};

        PipelineDynamicStateCreateInfo pdsci {
            .dynamicStateCount = 2,
            .pDynamicStates    = states,
        };

        PipelineShaderStageCreateInfo shaders[2] {
            {
                .stage  = ShaderStageFlagBits::eVertex,
                .module = *m_vert_shader,
                .pName  = "main",
            },
            {
                .stage  = ShaderStageFlagBits::eFragment,
                .module = *m_frag_shader,
                .pName  = "main",
            }};

        GraphicsPipelineCreateInfo gpci {
            .stageCount          = 2,
            .pStages             = shaders,
            .pVertexInputState   = &pcisci,
            .pInputAssemblyState = &piasci,
            .pViewportState      = &pvsci,
            .pRasterizationState = &prsci,
            .pMultisampleState   = &pmsci,
            .pDepthStencilState  = nullptr,
            .pColorBlendState    = &pcbsci,
            .pDynamicState       = &pdsci,
            .layout              = m_pipeline_layout.get(),
            .renderPass          = m_render_pass.get(),
            .subpass             = 0,
            .basePipelineHandle  = nullptr,
            .basePipelineIndex   = -1,
        };
        return m_device->createGraphicsPipeline(nullptr, gpci).value;
    }
    void VulkanTriangle::OnRecordRenderCommand(vk::CommandBuffer buffer,
                                               u32               currentFrame,
                                               u32               imageIdx,
                                               vk::Extent2D      screen_size)
    {
        using namespace vk;

        // reset buffer
        buffer.reset();

        CommandBufferBeginInfo cbbi {
            .flags = CommandBufferUsageFlagBits::eSimultaneousUse,
        };

        buffer.begin(cbbi);

        ClearValue clearValue {
            .color = {.float32 = std::array<f32, 4> {0.f, .0f, .0f, 1.f}}};

        RenderPassBeginInfo rpbi {
            .renderPass      = m_render_pass.get(),
            .framebuffer     = m_frame_buffers[imageIdx].get(),
            .renderArea      = {{0, 0}, screen_size},
            .clearValueCount = 1,
            .pClearValues    = &clearValue,
        };
        Viewport viewport {.x        = 0,
                           .y        = 0,
                           .width    = (f32)screen_size.width,
                           .height   = (f32)screen_size.height,
                           .minDepth = 0,
                           .maxDepth = 1.f};

        buffer.beginRenderPass(rpbi, SubpassContents::eInline);
        buffer.bindPipeline(PipelineBindPoint::eGraphics,
                             *m_render_pipeline);
        buffer.setViewport(0, viewport);
        buffer.bindVertexBuffers(0, m_vertex_buffer.get(), {0});
        buffer.bindIndexBuffer(m_index_buffer.get(), 0,
                                IndexType::eUint32);
        buffer.bindDescriptorSets(PipelineBindPoint::eGraphics, m_pipeline_layout.get(), 0, m_uniform_descriptor_sets[currentFrame], {});
        // buffer.draw(std::size(vertices), 1, 0, 0);
        buffer.drawIndexed(std::size(vertexIndices), 1, 0, 0, 0);
        buffer.endRenderPass();
        buffer.end();
    }

    vk::DescriptorSetLayout VulkanTriangle::OnCreateDescriptorSetLayout()
    {
        using namespace vk;
        DescriptorSetLayoutBinding mvpBinding{
            .binding = 0,
            .descriptorType = DescriptorType::eUniformBuffer,
            .descriptorCount = 1,
            .stageFlags = ShaderStageFlagBits::eVertex,
            .pImmutableSamplers = nullptr
        };
        DescriptorSetLayoutCreateInfo dslci{
            .flags = {},
            .bindingCount = 1,
            .pBindings = &mvpBinding,
        };

        return m_device->createDescriptorSetLayout(dslci);
    }

    void VulkanTriangle::OnUpdateBuffers(u32 imageIdx, u32 currentFrame) {
        using clock_type = std::chrono::high_resolution_clock;
        static clock_type::time_point begin_time = clock_type::now();

        UniformBufferObject mvp{};
        auto currentTime = clock_type::now();
        float time = std::chrono::duration<float, std::chrono::seconds::period>(currentTime - begin_time).count();

        mvp.model = glm::rotate(glm::mat4(1.0), time * glm::radians(90.f), glm::vec3(0.f,0.f,1.f));
        mvp.view = glm::lookAt(glm::vec3(2.f,2.f,2.f), glm::vec3(0.f,0.f,0.f), glm::vec3(0.f,0.f,1.f));
        mvp.proj = glm::perspective(glm::radians(45.f), (float)m_swapchian_size.width / (float)m_swapchian_size.height, 0.1f, 10.f);
        mvp.proj[1][1] *= -1; // inverse Y

        UploadBufferData(std::get<1>(m_uniform_buffers[currentFrame]).get(), &mvp, sizeof(mvp));
    }
}  // namespace core::app
